# -*- coding: utf-8 -*-

# FinBERT 微调

from datasets import Dataset
from transformers import (BertTokenizerFast, BertForSequenceClassification,
                          TrainingArguments, Trainer)
import torch, json, os

data = json.load(open("data/labeled.json"))
ds = Dataset.from_list([{"text": d["sentence"], "label": d["label"]+1} for d in data])  # 映射到 0/1/2

tok = BertTokenizerFast.from_pretrained("yiyanghkust/finbert-cn")
def tokenize(batch):
    return tok(batch["text"], padding="max_length", truncation=True, max_length=128)
ds = ds.map(tokenize, batched=True)

model = BertForSequenceClassification.from_pretrained(
            "yiyanghkust/finbert-cn", num_labels=3)

args = TrainingArguments(
    output_dir="models/finbert-risk",
    per_device_train_batch_size=16,
    learning_rate=2e-5,
    num_train_epochs=10,
    evaluation_strategy="no",
    save_total_limit=1,
    logging_steps=50
)

trainer = Trainer(model=model, args=args, train_dataset=ds)
trainer.train()
trainer.save_model("models/finbert-risk")
print("✅ 微调完成，已保存至 models/finbert-risk")

